package rewrite

import (
	"net/http"
	"regexp"
	"strings"
	"sync"
)

// Rule 定义重写规则
type Rule struct {
	Pattern     *regexp.Regexp
	Replacement string
	Type        string // "redirect" or "rewrite"
	Code        int    // HTTP状态码，用于重定向
}

// URLRewriter URL重写器
type URLRewriter struct {
	rules []*Rule
	mu    sync.RWMutex
}

func NewURLRewriter() *URLRewriter {
	return &URLRewriter{
		rules: make([]*Rule, 0),
	}
}

func (ur *URLRewriter) AddRule(pattern string, replacement string, ruleType string, code int) error {
	re, err := regexp.Compile(pattern)
	if err != nil {
		return err
	}

	ur.mu.Lock()
	defer ur.mu.Unlock()

	ur.rules = append(ur.rules, &Rule{
		Pattern:     re,
		Replacement: replacement,
		Type:        ruleType,
		Code:        code,
	})

	return nil
}

func (ur *URLRewriter) ClearRules() {
	ur.mu.Lock()
	defer ur.mu.Unlock()
	ur.rules = make([]*Rule, 0)
}

func (ur *URLRewriter) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ur.mu.RLock()
		rules := make([]*Rule, len(ur.rules))
		copy(rules, ur.rules)
		ur.mu.RUnlock()

		path := r.URL.Path
		var matched bool

		for _, rule := range rules {
			if rule.Pattern.MatchString(path) {
				newPath := rule.Pattern.ReplaceAllString(path, rule.Replacement)

				// 处理变量替换，如 $1, $2 等
				for i := 1; i <= 9; i++ {
					placeholder := "$" + string(rune('0'+i))
					if strings.Contains(newPath, placeholder) {
						matches := rule.Pattern.FindStringSubmatch(path)
						if len(matches) > i {
							newPath = strings.ReplaceAll(newPath, placeholder, matches[i])
						}
					}
				}

				if rule.Type == "redirect" {
					code := rule.Code
					if code == 0 {
						code = http.StatusMovedPermanently
					}
					http.Redirect(w, r, newPath, code)
					return
				}

				r.URL.Path = newPath
				matched = true
				break
			}
		}

		if !matched {
			next.ServeHTTP(w, r)
			return
		}

		next.ServeHTTP(w, r)
	})
}

// AddDefaultRules 添加常用的重写规则
func (ur *URLRewriter) AddDefaultRules() error {
	rules := []struct {
		pattern     string
		replacement string
		ruleType    string
		code        int
	}{
		// 规范化URL末尾的斜杠
		{`^(.+)/+$`, "$1", "rewrite", 0},
		// 将index.html重定向到目录
		{`^(.*)/index\.html$`, "$1/", "redirect", 301},
		// 规范化多个斜杠
		{`/{2,}`, "/", "rewrite", 0},
		// 添加www前缀的重定向
		{`^https?://(?!www\.)(.+)$`, "https://www.$1", "redirect", 301},
	}

	for _, rule := range rules {
		if err := ur.AddRule(rule.pattern, rule.replacement, rule.ruleType, rule.code); err != nil {
			return err
		}
	}

	return nil
}
